
"""
  States
"""

import roslib; roslib.load_manifest('sushi_smach')
import rospy
import smach
import smach_ros
import yaml
import tf
import geometry_msgs.msg
import actionlib

from collections import defaultdict
from geometry_msgs.msg import PoseStamped
from pr2_tasks.tasks import Tasks
from pr2_tasks.pickplace import PickPlace
from pr2_tasks.tableware_detection import TablewareDetection
from pr2_tasks.pickplace_definitions import PlaceGoal
from pr2_python import head,base
from pr2_python.world_interface import WorldInterface
from pr2_python.planning_scene_interface import get_planning_scene_interface
from pr2_python.arm_mover import ArmMover
from pr2_tasks.pickup_tableware import PickupTableware
from sushi_smach.GraspTurntable import GraspTurntable
import sushi_smach.msg


DIR = roslib.packages.get_pkg_dir('sushi_smach',required=True) + '/config/'
stream = file(DIR + 'poses.yaml')
poses = yaml.load(stream)

class Start(smach.State): 
  def __init__(self):
    smach.State.__init__(self,outcomes=['success'])

  def execute(self, userdata):
    rospy.loginfo('Executing Start State')

    return 'success'

# It maintains robot's object location knowledge
class WorldState:
  
  res = None

  def __init__(self):
    self.items = defaultdict(list)
  def add_item(self, object):
    rospy.loginfo("Add object into world state. Need to be worked on")

  def clear(self):
    self.items = defaultdict(list)


class CleanTable(smach.State): 
  def __init__(self):
    smach.State.__init__(self,outcomes=['success','failure'])
    self.cleanAction = actionlib.SimpleActionClient('/clean_table_action',sushi_smach.msg.CleanTableAction)

  def execute(self, userdata):
    rospy.loginfo('Executing Clean Table State')
    goal = sushi_smach.msg.CleanTableGoal()
    rospy.loginfo('Wait for Result')
    self.cleanAction.send_goal_and_wait(goal)
    res = self.cleanAction.get_result()
    rospy.loginfo('Result received = ' + str(res.result))

    if res.result == 'success':
      return 'success'
    else:
      return 'failure'

class ScanTable(smach.State): 
    
  world = None

  def __init__(self,world):
    self.world = world
    self.counter = 1
    smach.State.__init__(self,outcomes=['objects','no_object'])
    

  def execute(self, userdata):
    rospy.loginfo('Executing Scan Table')
    det = TablewareDetection()
    try:
      res = det.detect_objects()
      self.world.res = res
    except DetectionError:
      return 'no_object'


    if len(res.pickup_goals) >= 1:
      return 'objects'
    else:
      return 'no_object'

  def object_diameter(self, pg):
    y_max = pg.target.cluster.points[0].y
    y_min = pg.target.cluster.points[0].y
    
    for point in pg.target.cluster.points:
      if point.y > y_max:
        y_max = point.y
      if point.y < y_min:
        y_min = point.y
      return (y_max-y_min)

    
class NavigateTo(smach.State): 
  def __init__(self,world_state, base_controller, x, y, theta,hx,hy,ht):
    smach.State.__init__(self,outcomes=['success','failure'])
    self.x = x
    self.y = y
    self.theta = theta
    self.base = base_controller
    self.head = head.Head()
    self.hh = (hx,hy,ht)

  def execute(self, userdata):
    pose = (self.x, self.y, self.theta)
    rospy.loginfo('Executing Navigate To %s',str(pose))

    try:
      self.base.move_to(pose[0],pose[1],pose[2])
#self.base.move_manipulable_pose(pose[0], pose[1], pose[2])
      self.head.look_at_map_point(self.hh[0],self.hh[1],self.hh[2])
      print str(pose)
      print str(self.hh)
      return 'success'
    except:
      self.base.move_out_of_collision()
      
      try:
        self.base.move_to(pose[0],pose[1],pose[2])
        self.head.look_at_map_point(self.hh[0],self.hh[1],self.hh[2])
        return 'success'
      except:
        return 'failure'

class Celebrate(smach.State): 
  def __init__(self):
    smach.State.__init__(self,outcomes=['success','failure'])

  def execute(self, userdata):
    rospy.loginfo('Executing Celebration')

    return 'success'

class PickFromTurntable(smach.State):
  def __init__(self,world):
    smach.State.__init__(self,outcomes=['success','failure'])
    self.gt = GraspTurntable(arm_name='right_arm')

  def execute(self, userdata):
    self.gt.torso.down()

    try:
      if self.gt.pickup_object():
        self.gt.torso.up()
        return 'success'
      else:
        if self.gt.pickup_object():
          self.gt.torso.up()
          return 'success'
        else:
          return 'failure'
    except:
      try:
        if self.gt.pickup_object():
          self.gt.torso.up()
          return 'success'
        else:
          return 'failure'
      except:
        return 'failure'

class Pick(smach.State): 

  pickplace = None
  world = None
  base = None

  def __init__(self,world,base_controller):
    smach.State.__init__(self,outcomes=['success','failure'])
    self.pickplace = PickPlace()
    self.world = world
    self.base = base_controller

  def execute(self, userdata):
    rospy.loginfo('Executing Pick')
    res = self.world.res
    pickup_goals = []

    listener = tf.TransformListener()
    

#return 'success'
    
    for pg in res.pickup_goals:
      if pg.label  == 'cup':
        pickup_goals.append(pg)
#    for pg in res.pickup_goals:
#      if pg.label == 'small_cup':
#        pickup_goals.append(pg)
    for pg in res.pickup_goals:
      if pg.label == 'bowl':
        pickup_goals.append(pg)
    for pg in res.pickup_goals:
      if pg.label == 'small_bowl':
        pickup_goals.append(pg)
#    for pg in res.pickup_goals:
#      if pg.label == 'plate':
#        pickup_goals.append(pg)
    for pg in res.pickup_goals:
      if pg.label == 'graspable':
        pickup_goals.append(pg)
    
#pickup_goals = pickup_goals[0:2]
    self.world.res.pickup_goals = pickup_goals

    x = 0
    y = 0
    z = 0

#arm_mover = ArmMover()

    if pickup_goals[0].label is 'graspable':
      for point in pickup_goals[0].target.cluster.points:
        x += point.x / len(pickup_goals[0].target.cluster.points)
        y += point.y / len(pickup_goals[0].target.cluster.points)
        z += point.z / len(pickup_goals[0].target.cluster.points)
      if y < 0:
        pickup_goals[0].arm_name = 'right_name'
        if len(pickup_goals) > 1:
          pickup_goals[1].arm_name = 'left_arm'
      else:
        pickup_goals[1].arm_name = 'left_arm'
        if len(pickup_goals) > 1:
          pickup_goals[1].arm_name = 'right_arm'
    else:
      if pickup_goals[0].target.potential_models[0].pose.pose.position.y < 0:
        pickup_goals[0].arm_name = 'right_arm'
        if len(pickup_goals) > 1:
          pickup_goals[1].arm_name = 'left_arm'
      else:
        pickup_goals[0].arm_name = 'left_arm'
        if len(pickup_goals) > 1:
          pickup_goals[1].arm_name = 'right_arm'
      x = pickup_goals[0].target.potential_models[0].pose.pose.position.x
      y = pickup_goals[0].target.potential_models[0].pose.pose.position.y
      z = pickup_goals[0].target.potential_models[0].pose.pose.position.z
      print str(pickup_goals[0].target.potential_models[0].pose)
    goal_pose = PoseStamped()
    goal_pose.pose.position.x = x
    goal_pose.pose.position.y = y + 0.15
    goal_pose.pose.position.z = z
    goal_pose.pose.orientation.w = 1.0
    goal_pose.header.frame_id = '/torso_lift_link'


    try:
      print x
      print y 
      print z
      arm_name = None
      if pickup_goals[0].arm_name is 'left_arm':
        arm_name = 'left_arm'
      else:
        arm_name = 'right_arm'

      print str(arm_name)

      (trans,rot) = listener.lookupTransform('/base_link','/map',rospy.Time.now())
      x = trans[0] + x
      y = trans[1] + y
      z = trans[2] + z
       
      self.base.move_manipulable_pose(x, y, z)
      self.pickplace.pickup(pickup_goals[0])
#handle = arm_mover.move_to_goal(pickup_goals[0].arm_name, goal_pose) 
    except :
      print str(e)
      return 'failure'

    if len(pickup_goals) == 1:
      return 'success'

    if pickup_goals[1].label is 'graspable':
      for point in pickup_goals[1].target.cluster.points:
        x = point.x / len(pickup_goals[1].target.cluster.points)
        y = point.y / len(pickup_goals[1].target.cluster.points)
        z = point.z / len(pickup_goals[1].target.cluster.points)
    else:
      x = pickup_goals[1].target.potential_models[0].pose.pose.position.x
      y = pickup_goals[1].target.potential_models[0].pose.pose.position.y
      z = pickup_goals[1].target.potential_models[0].pose.pose.position.z

    goal_pose = PoseStamped()
    goal_pose.pose.position.x = x   
    goal_pose.pose.position.y = y + 0.15
    goal_pose.pose.position.z = z
    goal_pose.pose.orientation.w = 1.0
    goal_pose.header.frame_id = '/torso_lift_link'

    try:
#self.base.move_manipulable_pose(x, y, z, left=(pickup_goals[1].arm_name is 'left_arm'))
      self.pickplace.pickup(pickup_goals[1])
#handle = arm_mover.move_to_goal(pickup_goals[1].arm_name, goal_pose)

    except:
      return 'failure'

    return 'success'

class Putdown(smach.State): 
  pickplace = None
  world = None
  def __init__(self,world):
    smach.State.__init__(self,outcomes=['success','failure'])
    self.pickplace = PickPlace()
    self.world = world
    
  def execute(self, userdata):
    rospy.loginfo('Executing Putdown')
    place_pose = self.getPose()
    pg = self.world.res.pickup_goals(0)
    place_goal = PlaceGoal(pg.arm_name, place_pose,
                         collision_support_surface_name = pg.collision_support_surface_name,
                         collision_object_name = pg.collision_object_name)

    #and put it back down
    try:
      pickplace.place(place_goal)

    except:
      return 'failure'

    if len(self.items) == 1:
      return 'success'

    pg = self.world.res.pickup_goals(1)
    place_pose.pose.pose.position.y += 0.1
    place_goal = PlaceGoal(pg.arm_name, place_pose,
                          collision_support_surface_name = pg.collision_support_surface_name,
                          collision_object_name = pg.collision_object_name)
    try:
      pickplace.place(place_goal)
    except:
      return 'failure'
        
    return 'success'        
  
#def getPose():

class MovePutdown(smach.State): 
  def __init__(self,world,nav_states):
    smach.State.__init__(self,outcomes=['success','failure'])
    self.tasks= Tasks()

  def execute(self, userdata):
    rospy.loginfo('Executing Move Putdown')
    place_stamped = PoseStamped()
    place_stamped.header.frame_id = '/map'
    place_stamped.pose.position.x = 0.462
    place_stamped.pose.position.y = 1.589
    place_stamped.pose.position.z = 0.752
    place_stamped.pose.orientation.w = 1.0
   
    rospy.loginfo('Going to place!')
    try:
      self.tasks.go_and_place('right_arm', place_stamped)
      return 'success'
    except:
      return 'failure'

class TablePickup(smach.State): 
  def __init__(self,world,nav_states):
    smach.State.__init__(self,outcomes=['success','failure'])
    self.pickplace = PickPlace()
    self.world = world

  def execute(self, userdata):
    rospy.loginfo('Executing Table Pickup')
    pickup_goals = []
    res = self.world.res
    
    for pg in res.pickup_goals:
     if pg.label == 'large_cup' and ('large_cup' not in self.items or 'small_cup' not in self.items):
       pickup_goals.append(pg)
    for pg in res.pickup_goals:
      if pg.label == 'small_cup' and ('large_cup' not in self.items or 'small_cup' not in self.items):
        pickup_goals.append(pg)
    for pg in res.pickup_goals:
      if pg.label == 'large_bowl' and ('large_bowl' not in self.items or 'small_bowl' not in self.items):
        pickup_goals.append(pg)
    for pg in res.pickup_goals:
      if pg.label == 'small_bowl' and ('large_bowl' not in self.items or 'small_bowl' not in self.items):
        pickup_goals.append(pg)
    for pg in res.pickup_goals:
      if pg.label == 'plate' and 'plate' not in self.items:
        pickup_goals.append(pg)
    
    pickup_goals = pickup_goals[0:2]
    self.world.res.pickup_goals = pickup_goals
    if pickup_goals[0].target.potential_models(0).pose.pose.y > 0:
      pickup_goals[0].arm_name = 'right_arm'
      pickup_goals[1].arm_name = 'left_arm'
    else:
      pickup_goals[0].arm_name = 'left_arm'
      pickup_goals[1].arm_name = 'right_arm'

    pickplace = PickPlace()
    x = pickup_goals[0].target.potential_models(0).pose.pose.position.x
    y = pickup_goals[0].target.potential_models(0).pose.pose.position.y
    z = pickup_goals[0].target.potential_models(0).pose.pose.position.z
    try:
      move_manipulable_pose(x, y, z, left=(pickup_goals(0).arm_name is 'left_arm'))
      self.pickplace.pickup(pickup_goals(0))
    except:
      return 'failure'

    x = pickup_goals[1].target.potential_models(0).pose.pose.position.x
    y = pickup_goals[1].target.potential_models(0).pose.pose.position.y
    z = pickup_goals[1].target.potential_models(0).pose.pose.position.z

    try:
      move_manipulable_pose(x, y, z, left=(pickup_goals(1).arm_name is 'left_arm'))
      self.pickplace.pickup(pickup_goals(1))
    except:
      return 'failure'

    return 'success'
class TablePutdown(smach.State): 
  def __init__(self,world,nav_states):
    smach.State.__init__(self,outcomes=['success','failure'])

  def execute(self, userdata):
    rospy.loginfo('Executing Table Putdown')
    place_pose = self.getPose()
    if len(self.world.res.pickup_goals) >= 1:
      pg = self.world.res.pickup_goals(0)
      place_goal = PlaceGoal(pg.arm_name, place_pose,
                           collision_support_surface_name = pg.collision_support_surface_name,
                           collision_object_name = pg.collision_object_name)

    #and put it back down
      try:
        pickplace.place(place_goal)
        self.items.append(pg.label)

      except:
        return 'failure'
      
      if len(self.items) >= 2:
        pg = self.world.res.pickup_goals(1)
        place_pose.pose.position.y += 0.1
        place_goal = PlaceGoal(pg.arm_name, place_pose,
                           collision_support_surface_name = pg.collision_support_surface_name,
                           collision_object_name = pg.collision_object_name)
        try:
          pickplace.place(place_goal)
          self.items.append(pg.label)
        except:
          return 'failure'
        
      else:
        return 'success'        
    else:
      return 'failure'
  
  def getPose():
    return PoseStamped()

    return 'success'

class CheckStatus(smach.State): 
  def __init__(self):
    smach.State.__init__(self,outcomes=['success','failure'])

  def execute(self, userdata):
    rospy.loginfo('Executing Check Status')
    return 'success'
